#include <iostream>
#include <thread>
#include <vector>
#include <queue>
#include <mutex>
#include <condition_variable>
#include <functional>
#include <atomic>
#include <future>
class ThreadPool
{
private:
    // 线程池中的线程
    std::vector<std::thread> threads;
    // 任务队列
    std::queue<std::function<void()>> tasks;
    // 互斥锁
    std::mutex mtx;
    // 条件变量
    std::condition_variable cv;
    // 是否关闭线程池
    std::atomic<bool> stop;
    // 线程函数
    void work()
    {
        while (true)
        {
            std::function<void()> task;
            {
                std::unique_lock<std::mutex> lock(mtx);
                cv.wait(lock, [this]
                        { return stop || !tasks.empty(); });
                if (stop && tasks.empty())
                {
                    return;
                }
                task = tasks.front();
                tasks.pop();
            }
            task();
        }
    }

public:
    // 构造函数
    ThreadPool(int threadNum) : stop(false)
    {
        // 创建threadNum个线程
        for (int i = 0; i < threadNum; ++i)
        {
            // 创建线程并将线程函数设置为work函数
            threads.emplace_back(std::thread(std::bind(&ThreadPool::work, this)));
        }
    };
    // 析构函数
    ~ThreadPool()
    {
        // 关闭线程池
        stop = true; // 原子操作，不需要加锁
        cv.notify_all();
        // 等待所有线程执行完毕
        for (std::thread &t : threads)
        {
            t.join();
        }
    }
    // 添加任务
    template <class F, class... Args>
    void enqueue(F &&f, Args &&...args)
    {
        // 将任务函数和参数绑定
        auto task = std::bind(std::forward<F>(f), std::forward<Args>(args)...);
        {
            std::unique_lock<std::mutex> lock(mtx);
            // 不允许在线程池停止后添加任务
            if (stop)
            {
                throw std::runtime_error("enqueue on stopped ThreadPool");
            }
            tasks.emplace(task);
        }
        cv.notify_one();
    }
    template <class F, class... Args>
    // typename 关键字用于告诉编译器 std::result_of<F(Args...)>::type 是一个类型。
    // std::result_of<F(Args...)>::type 用于推导可调用对象 f 在传入 Args... 参数时的返回类型
    // std::future 则用于异步获取该返回值
    auto addTask(F &&f, Args &&...args) -> std::future<typename std::result_of<F(Args...)>::type>
    {
        using returnType = typename std::result_of<F(Args...)>::type;
        // 将任务函数和参数绑定
        // std::packaged_task 是一个类模板，用于封装一个可调用对象，并将其结果与一个 std::future 对象关联起来
        auto task = std::make_shared<std::packaged_task<returnType()>>(
            std::bind(std::forward<F>(f), std::forward<Args>(args)...));
        // 获取 std::future 对象
        std::future<returnType> res = task->get_future();
        {
            std::unique_lock<std::mutex> lock(mtx);
            if (stop)
            {
                throw std::runtime_error("enqueue on stopped ThreadPool");
            }
            // 使用 Lambda 表达式的原因 解决类型不匹配
            tasks.emplace([task]()
                          { (*task)(); });
        }
        cv.notify_one();
        return res;
    }
};

// 示例任务函数
void exampleTask(int id)
{
    std::cout << "Task " << id << " is running on thread " << std::this_thread::get_id() << std::endl;
    std::this_thread::sleep_for(std::chrono::seconds(1));
    std::cout << "Task " << id << " is finished." << std::endl;
}

int main()
{
    // 创建一个包含 4 个线程的线程池
    ThreadPool pool(8);

    // // 向线程池中添加 8 个任务
    for (int i = 10; i < 18; ++i)
    {
        pool.enqueue(exampleTask, i);
    }

    // 向线程池添加一些任务
    std::vector<std::future<void>> futures;
    for (int i = 0; i < 8; ++i)
    {
        futures.emplace_back(pool.addTask(exampleTask, i));
    }

    // 等待所有任务完成
    for (auto &future : futures)
    {
        future.wait();
    }

    return 0;
}